package main

import (
	"encoding/json"
	"errors"
	"github.com/dgraph-io/badger/v4"
	"log"
	"net/http"
)

type KeyValue struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

type Response struct {
	Success bool   `json:"success"`
	Result  any    `json:"result"`
	Message string `json:"message"`
}

func getSuccessResult(resp any, message string) *Response {
	return &Response{
		Success: true,
		Result:  resp,
		Message: message,
	}
}

func writeHandler(db *badger.DB, config *Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		var req KeyValue
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		group := getGroup(r, config)
		fullKey := getKeyByte(group, req.Key, config)

		wb := db.NewWriteBatch()
		defer wb.Cancel()

		if err := wb.Set(fullKey, []byte(req.Value)); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if err := wb.Flush(); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		w.WriteHeader(http.StatusOK)
		err := json.NewEncoder(w).Encode(getSuccessResult(nil, ""))
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
		}
	}
}

func getHandler(db *badger.DB, config *Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		key := r.URL.Query().Get("key")
		group := getGroup(r, config)
		fullKey := getKeyByte(group, key, config)

		var value []byte
		err := db.View(func(txn *badger.Txn) error {
			item, err := txn.Get(fullKey)
			if err != nil {
				if errors.Is(err, badger.ErrKeyNotFound) {
					json.NewEncoder(w).Encode(getSuccessResult(nil, "key not exist"))
					return nil
				} else {
					return err
				}
			}
			value, err = item.ValueCopy(nil)
			return err
		})

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if value != nil {
			json.NewEncoder(w).Encode(getSuccessResult(string(value), ""))
		}
	}
}

func deleteHandler(db *badger.DB, config *Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		key := r.URL.Query().Get("key")
		group := getGroup(r, config)

		fullKey := getKeyByte(group, key, config)

		wb := db.NewWriteBatch()
		defer wb.Cancel()

		if err := wb.Delete(fullKey); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if err := wb.Flush(); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode(getSuccessResult(nil, ""))
	}
}

func batchWriteHandler(db *badger.DB, config *Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		var kvList []KeyValue
		err := json.NewDecoder(r.Body).Decode(&kvList)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		group := getGroup(r, config)

		wb := db.NewWriteBatch()
		defer wb.Cancel()

		for _, kv := range kvList {
			fullKey := getKeyByte(group, kv.Key, config)
			value := []byte(kv.Value)
			if err := wb.Set(fullKey, value); err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
		}

		if err := wb.Flush(); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode(getSuccessResult(nil, ""))
	}
}

func batchGetHandler(db *badger.DB, config *Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		var keys []string
		if err := json.NewDecoder(r.Body).Decode(&keys); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		group := getGroup(r, config)

		values := make(map[string]string)
		err := db.View(func(txn *badger.Txn) error {
			for _, key := range keys {
				fullKey := getKeyByte(group, key, config)
				item, err := txn.Get(fullKey)
				if err != nil {
					if errors.Is(err, badger.ErrKeyNotFound) {
						continue
					}
					return err
				}
				valueCopy, err := item.ValueCopy(nil)
				if err != nil {
					return err
				}
				values[key] = string(valueCopy)
			}
			return nil
		})

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Header().Set("Content-Type", "application/json; charset=UTF-8")
			json.NewEncoder(w).Encode(getSuccessResult(values, ""))
		}
	}
}

func getByPrefixHandler(db *badger.DB, config *Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		group := getGroup(r, config)
		key := r.URL.Query().Get("prefix")
		fullKey := getKeyByte(group, key, config)

		values := make(map[string]string)
		err := db.View(func(txn *badger.Txn) error {
			opts := badger.DefaultIteratorOptions
			opts.PrefetchSize = 20
			it := txn.NewIterator(opts)
			defer it.Close()
			for it.Seek(fullKey); it.ValidForPrefix(fullKey); it.Next() {
				item := it.Item()
				k := item.Key()
				err := item.Value(func(v []byte) error {
					k = k[len(group)+1:]
					values[string(k)] = string(v)
					return nil
				})
				if err != nil {
					return err
				}
			}
			return nil
		})
		if err != nil {
			log.Println("Failed to iterator keys and values from the cache.", "error", err)
		}
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		json.NewEncoder(w).Encode(getSuccessResult(values, ""))
	}
}

func batchDeleteHandler(db *badger.DB, config *Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		var keys []string
		if err := json.NewDecoder(r.Body).Decode(&keys); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		group := getGroup(r, config)

		txn := db.NewTransaction(true)
		defer txn.Discard()

		for _, key := range keys {
			fullKey := getKeyByte(group, key, config)
			if err := txn.Delete(fullKey); err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
		}

		if err := txn.Commit(); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode(getSuccessResult(nil, ""))
	}
}

func getKeyByte(group []byte, key string, config *Config) []byte {
	totalLength := len(group) + len(key) + 1
	// 预先分配足够的空间
	combined := make([]byte, 0, totalLength)
	// 使用 copy 将每个切片复制到新的切片中
	combined = append(group, config.Split)
	combined = append(combined, []byte(key)...)
	return combined
}

func getGroup(r *http.Request, config *Config) []byte {
	group := r.URL.Query().Get("group")
	if group == "" {
		return []byte(config.DefaultGroup)
	}
	return []byte(group)
}
